import os
import subprocess
from datetime import datetime
import pymysql
import argparse
import configparser
import sys

# 默认设置
默认操作模式 = 'backup'
默认备份目录 = os.path.expanduser("~/backup/")
配置文件路径 = os.path.join(os.path.dirname(__file__), 'config.ini')

def 加载配置():
    if not os.path.exists(配置文件路径):
        print(f"错误：配置文件不存在：{配置文件路径}")
        sys.exit(1)
    
    配置解析器 = configparser.ConfigParser()
    配置解析器.read(配置文件路径)
    
    配置 = {
        'backup': dict(配置解析器['backup']),
        'restore': dict(配置解析器['restore'])
    }
    
    # 将端口号转换为整数
    for section in ['backup', 'restore']:
        if 'port' in 配置[section]:
            配置[section]['port'] = int(配置[section]['port'])
    
    return 配置

def 获取数据库列表(游标):
    """获取所有数据库（排除系统数据库）"""
    游标.execute("SHOW DATABASES;")
    数据库列表 = 游标.fetchall()
    return [数据库[0] for 数据库 in 数据库列表 if 数据库[0] not in ['information_schema', 'performance_schema', 'mysql', 'sys']]

def 执行备份(数据库配置, 数据库列表, 备份文件路径):
    """执行数据库备份"""
    数据库参数 = ' '.join([f'{数据库}' for 数据库 in 数据库列表])
    mysqldump命令 = (
        f"mysqldump -u {数据库配置['user']} -p{数据库配置['password']} "
        f"-h {数据库配置['host']} -P {数据库配置['port']} "
        f"--single-transaction --quick --lock-tables=false "
        f"--set-gtid-purged=OFF --databases {数据库参数} > {备份文件路径}"
    )
    进程 = subprocess.run(mysqldump命令, shell=True, stderr=subprocess.PIPE)
    if 进程.returncode != 0:
        print(f"备份错误: {进程.stderr.decode()}")
        return False
    return True

def 执行恢复(数据库配置, 恢复文件):
    """执行数据库恢复"""
    mysql命令 = (
        f"mysql -u {数据库配置['user']} -p{数据库配置['password']} "
        f"-h {数据库配置['host']} -P {数据库配置['port']}"
    )
    
    进程 = subprocess.run(
        f"{mysql命令} < {恢复文件}",
        shell=True,
        stderr=subprocess.PIPE
    )
    
    if 进程.returncode != 0:
        print(f"恢复错误: {进程.stderr.decode()}")
        return False
    return True

def 主函数():
    解析器 = argparse.ArgumentParser(description='MySQL 数据库备份和恢复脚本')
    解析器.add_argument('操作模式', choices=['backup', 'restore'], default=默认操作模式, nargs='?',
                      help='操作模式: "backup"(备份) 或 "restore"(恢复)')
    解析器.add_argument('--backup-file', dest='指定备份文件',
                      help='指定用于恢复的备份文件路径')
    解析器.add_argument('--config', dest='配置文件',
                      default=配置文件路径,
                      help='指定配置文件路径')
    
    参数 = 解析器.parse_args()
    
    try:
        配置信息 = 加载配置()
        
        if 参数.操作模式 == 'backup':
            备份文件名 = f"{datetime.now().strftime('%Y%m%d%H%M%S')}_full_backup_exclude_sys.sql"
            备份文件路径 = os.path.join(默认备份目录, 备份文件名)
            
            if not os.path.exists(默认备份目录):
                os.makedirs(默认备份目录)
            
            连接 = pymysql.connect(**配置信息['backup'])
            游标 = 连接.cursor()
            
            print("正在获取数据库列表...")
            待备份数据库列表 = 获取数据库列表(游标)
            print(f"备份目标: {', '.join(待备份数据库列表)}")
            
            print("开始备份...")
            备份成功 = 执行备份(配置信息['backup'], 待备份数据库列表, 备份文件路径)
            
            if 备份成功:
                print(f"备份完成。保存位置: {备份文件路径}")
            else:
                print("备份失败。")
            
            游标.close()
            连接.close()
        
        elif 参数.操作模式 == 'restore':
            指定备份文件 = 参数.指定备份文件
            if not 指定备份文件:
                备份文件列表 = [文件 for 文件 in os.listdir(默认备份目录) if 文件.endswith('.sql')]
                if not 备份文件列表:
                    print("错误：没有找到备份文件。")
                    sys.exit(1)
                
                指定备份文件 = os.path.join(
                    默认备份目录,
                    max(备份文件列表, key=lambda x: os.path.getctime(os.path.join(默认备份目录, x)))
                )
            
            if os.path.exists(指定备份文件):
                print(f"开始恢复: {指定备份文件}")
                恢复成功 = 执行恢复(配置信息['restore'], 指定备份文件)
                
                if 恢复成功:
                    print("恢复完成。")
                else:
                    print("恢复失败。")
            else:
                print(f"错误：指定的备份文件不存在: {指定备份文件}")
    
    except Exception as e:
        print(f"发生错误: {e}")
        sys.exit(1)

if __name__ == "__main__":
    主函数()